#include "network.h"
#include "logging.h"
#include "owl_math.h"
#include "connected_layer.h"

int GetCurrentBatch(Network net) {
    int batch_num = (*net.seen) / (net.batch * net.subdivisions);
    return batch_num;
}

void ResetMomentum(Network net) {
    if (net.momentum == 0) {
        return;
    }
    net.learing_rate = 0.0;
    net.momentum = 0.0;
    net.decay = 0.0;
}

Network MakeNetwork(int n) {
    Network net = {0};
    net.layer_count = n;
    net.layers = (Layer *)OwlAlloc(net.layer_count, sizeof(Layer));
    net.seen = (int *)OwlAlloc(1, sizeof(int));
    return net;
}

void ForwardNetwork(Network net, NetworkState state) {
    LOG(INFO) << "Network forward";
    state.data_workspace = net.data_workspace;
    int i = 0;
    for (i = 0; i < net.layer_count; ++i) {
        state.index = i;
        Layer layer = net.layers[i];
        LOG(INFO) << layer.type << " input: " << layer.inputs << " output: " << layer.outputs;
        if (layer.data_delta) {
            OwlScal(layer.outputs * layer.batch, 0, layer.data_delta, 1);
        }

        layer.Forward(layer, state);
        state.data_input = layer.data_output;
    }

}

void BackwardNetwork(Network net, NetworkState state) {
    Float *original_input = state.data_input;
    Float *original_delta = state.data_delta;
    state.data_workspace = net.data_workspace;
    int i = 0;
    for (i = net.layer_count - 1; i >= 0; --i) {
        state.index = i;
        if (i == 0) {
            state.data_input = original_input;
            state.data_delta = original_delta;
        } else {
            Layer prev = net.layers[i - 1];
            state.data_input = prev.data_output;
            state.data_delta = prev.data_delta;
        }

        Layer layer = net.layers[i];
        LOG(INFO) << "Backward layer type: " << layer.type;
        layer.Backward(layer, state);

    }
}

void UpdateNetwork(Network net, NetworkState state) {
    int i;
    int update_batch = net.batch * net.subdivisions;
    Float rate = 0.1;
    for (i = 0; i < net.layer_count; ++i) {
        Layer layer = net.layers[i];
        if (layer.Update) {
            layer.Update(layer, update_batch, rate, net.momentum, net.decay);
        }
    }
}

Float TrainNetworkDatum(Network net, Float *x, Float *y) {
    NetworkState state;
    *(net.seen) += net.batch;
    state.index = 0;
    state.data_delta = 0;
    state.data_input = x;
    state.data_truth = y;
    state.train = 1;
    OWLLOG("TrainNetworkDatum", x, 2);
    LOG(INFO) << "begin to forward";
    ForwardNetwork(net, state);
    BackwardNetwork(net, state);

}
